Skip to main content

Quick Start

Setup

  • Please refer to the installation guide to setup your environment first

  • Install finetune dependencies:

    pip install "cogkit[finetune]@git+https://github.com/THUDM/CogKit.git"
  • We provide various training scripts and example datasets in the CogKit/quickstart directory. Please clone the repository before training:

    git clone https://github.com/THUDM/CogKit.git

Data

Before fine-tuning, you need to prepare your dataset according to the expected format. See the data format documentation for details on how to structure your data

Training

info

We recommend that you read the corresponding model card before starting training to follow the parameter settings requirements and fine-tuning best practices

  1. Navigate to the CogKit/ directory after cloning the repository

    cd CogKit/
  2. Choose the appropriate subdirectory from the quickstart/scripts based on your task type and distribution strategy. For example, t2i corresponds to text-to-image task

  3. Review and adjust the parameters in config.yaml in the selected training directory

  4. Run the script in the selected directory:

    bash start_train.sh

Load Fine-tuned Model

Merge Checkpoint

After fine-tuning, you need to use the merge.py script to merge the distributed checkpoint weights into a single checkpoint (except for QLoRA fine-tuning). The script can be found in the quickstart/tools/converters directory. For example:

cd quickstart/tools/converters
python merge.py --checkpoint_dir ckpt/ --output_dir output_dir/
# Add --lora option if you are using LoRA fine-tuning

Load Checkpoint

You can pass the output_dir to the --lora_model_id_or_path option if you are using LoRA fine-tuning, or to the --transformer_path option if you are using FSDP fine-tuning. For more details, please refer to the inference guide.